#!/usr/bin/env python3
# encoding: utf-8
# @Author: Zecous
# @Date: 2025/07/18
import cv2
import numpy as np
import rospy
import signal
import time
import math
from sensor_msgs.msg import Image
from std_msgs.msg import String, Bool
from std_srvs.srv import Trigger, TriggerResponse
from cv_bridge import CvBridge
from ainex_sdk import misc, common
from ainex_example.color_common import Common

class HurdlesTracker(Common):
    def __init__(self, name):
        # 初始化头部位置
        self.head_pan_init = 500   # 左右舵机的初始值
        self.head_tilt_look_down = 280  # 低头位置
        
        # 初始化ROS节点
        rospy.init_node(name)
        self.name = name
        self.running = True
        self.active = False  # 默认不激活，等待外部启动
        
        super().__init__(name, self.head_pan_init, self.head_tilt_look_down)
        
        # 设置退出处理函数
        signal.signal(signal.SIGINT, self.shutdown)
        
        # Initialize CV Bridge
        self.bridge = CvBridge()
        
        # State variables 
        self.current_image = None
        self.current_state = "approaching"  # approaching -> ready -> done
        self.hurdles_executed = False
        self.task_completed = False
        # 目标线稳定计数器
        self.target_y_counter = 0
        
        # 对齐控制参数
        self.position_gain = 1.0     # 位置控制增益
        self.angle_gain = 1.2        # 角度控制增益
        self.small_step_ratio = 0.6  # 小步前进比例
        
        # 从ROS参数服务器加载配置
        self.load_config()
        
        # 初始化运动参数
        self.init_movement_params()
        
        # 图像参数
        self.image_width = rospy.get_param('~image_width', 640)
        self.image_height = rospy.get_param('~image_height', 480)
        self.center_x = self.image_width // 2
        
        # 从配置加载位置判断参数
        self.load_position_params()
        
        # ROS通信接口
        self.setup_ros_interface()

        # 显示窗口控制
        self.update_display_config()
        
        # Subscribe to camera topic
        self.image_subscriber = rospy.Subscriber('/camera/image_raw', Image, self.image_callback)
        
        # 初始化机器人动作
        self.motion_manager.run_action('walk_ready')
        
        # 设置低头位置
        self.set_servo(self.head_pan_init, self.head_tilt_look_down)
        time.sleep(2)  # 等待舵机到位
        
        common.loginfo('Hurdles tracker initialized with unified control')
        print("Hurdles task ready, waiting for activation...")

    def update_display_config(self):
        """从参数服务器更新显示配置"""
        display_ns = '/system/display'
        main_flag = rospy.get_param(f'{display_ns}/show_hurdles_window', None)
        if main_flag is None:
            main_flag = rospy.get_param(f'{display_ns}/show_main_window', True)
        self.display_main_window = bool(main_flag)

        mask_flag = rospy.get_param(f'{display_ns}/show_detection_masks', True)
        self.display_mask_window = self.display_main_window and bool(mask_flag)

    @staticmethod
    def _safe_destroy_window(window_name):
        try:
            cv2.destroyWindow(window_name)
        except cv2.error:
            pass
        except Exception:
            pass

    def setup_ros_interface(self):
        """设置ROS通信接口"""
        # 状态发布器
        self.status_pub = rospy.Publisher('/hurdles/status', String, queue_size=1)
        self.active_pub = rospy.Publisher('/hurdles/active', Bool, queue_size=1)
        self.progress_pub = rospy.Publisher('/hurdles/progress', String, queue_size=1)
        
        # 控制服务
        self.start_service = rospy.Service('/hurdles/start', Trigger, self.start_service_callback)
        self.stop_service = rospy.Service('/hurdles/stop', Trigger, self.stop_service_callback)
        self.reset_service = rospy.Service('/hurdles/reset', Trigger, self.reset_service_callback)
        
        # 控制命令订阅
        self.control_sub = rospy.Subscriber('/hurdles/control', String, self.control_callback)

    def start_service_callback(self, req):
        """启动服务回调"""
        self.active = True
        self.running = True
        self.task_completed = False
        self.hurdles_executed = False
        self.current_state = "approaching"
        self.publish_status("started")
        return TriggerResponse(success=True, message="Hurdles task started")

    def stop_service_callback(self, req):
        """停止服务回调"""
        self.active = False
        self.gait_manager.stop()
        self.publish_status("stopped")
        return TriggerResponse(success=True, message="Hurdles task stopped")

    def reset_service_callback(self, req):
        """重置服务回调"""
        self.active = False
        self.gait_manager.stop()
        self.current_state = "approaching"
        self.hurdles_executed = False
        self.task_completed = False
        self.publish_status("reset")
        return TriggerResponse(success=True, message="Hurdles task reset")

    def control_callback(self, msg):
        """控制命令回调"""
        command = msg.data.lower()
        if command == "start":
            self.active = True
            self.running = True
            self.task_completed = False
        elif command == "stop":
            self.active = False
            self.gait_manager.stop()
        elif command == "reset":
            self.active = False
            self.gait_manager.stop()
            self.current_state = "approaching"
            self.hurdles_executed = False
            self.task_completed = False

    def publish_status(self, status):
        """发布状态"""
        msg = String()
        msg.data = status
        self.status_pub.publish(msg)
        
        active_msg = Bool()
        active_msg.data = self.active
        self.active_pub.publish(active_msg)
        
        progress_msg = String()
        progress_msg.data = f"State: {self.current_state}, Executed: {self.hurdles_executed}, Completed: {self.task_completed}"
        self.progress_pub.publish(progress_msg)



    def load_config(self):
        """从ROS参数服务器加载配置"""
        try:
            # 从参数服务器获取颜色配置
            blue_config = rospy.get_param('/image_processing/colors/blue', {})
            self.blue_lower = np.array(blue_config.get('lower', [0, 94, 0]))
            self.blue_upper = np.array(blue_config.get('upper', [243, 179, 113]))
            
            # 从参数服务器获取舵机配置
            servo_config = rospy.get_param('/servo', {})
            self.head_pan_init = servo_config.get('head_pan_init', 500)
            self.head_tilt_look_down = servo_config.get('head_tilt_look_down', 280)
            
            # print(f"配置加载成功 - 蓝色范围: {self.blue_lower} ~ {self.blue_upper}")
            
        except Exception as e:
            print(f"配置加载失败，使用默认值: {e}")
            # 使用默认配置
            self.blue_lower = np.array([0, 94, 0])
            self.blue_upper = np.array([243, 179, 113])
            self.head_pan_init = 500
            self.head_tilt_look_down = 280

    def load_position_params(self):
        """从配置加载位置判断参数"""
        try:
            # 从参数服务器获取跨栏任务配置
            hurdles_config = rospy.get_param('/tasks/hurdles', {})
            position_ratios = hurdles_config.get('position_ratios', {})
            
            # 加载位置比例参数
            target_y_ratio = position_ratios.get('target_y_ratio', 0.75)
            self.target_y = int(self.image_height * target_y_ratio)
            # 需要连续达到目标线才认为到位（避免抖动导致误判）
            self.target_y_stable_frames = int(position_ratios.get('target_y_stable_frames', 2))
            # 重置计数器
            self.target_y_counter = 0
            
            print(f"位置参数加载成功 - hurdles_config: {hurdles_config}")
            print(f"位置参数加载成功 - position_ratios: {position_ratios}")
            print(f"位置参数加载成功 - 目标Y位置: {self.target_y} (比例: {target_y_ratio})")
            
        except Exception as e:
            print(f"位置参数加载失败，使用默认值: {e}")
            # 使用默认配置
            self.target_y = int(self.image_height * 0.75)

    def init_movement_params(self):
        """初始化运动参数"""
        self.calib_config = common.get_yaml_data('/home/ubuntu/ros_ws/src/ainex_example/config/calib.yaml')
        
        # 获取步态参数
        self.walking_param = self.gait_manager.get_gait_param()
        self.walking_param['step_height'] = 0.02    # 步高
        self.walking_param['pelvis_offset'] = 5      # 盆骨偏移
        self.walking_param['hip_pitch_offset'] = 15  # 髋关节偏移
        self.walking_param['body_height'] = 0.035    # 身体高度
        self.walking_param['init_y_offset'] = 0.000    # 初始y偏移
        
        # 运动参数设置（参考循线节点）
        self.dsp_param = [350, 0.2, 0.025]  # DSP参数
        self.arm_swap = 60                   # 摆臂幅度（参考循线节点）
        self.step_distance = 0.008          # 前进步长

    def image_callback(self, msg):
        """ROS图像回调函数"""
        try:
            self.current_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except Exception as e:
            rospy.logerr(f"Failed to convert image: {e}")

    def detect_blue_line(self, image):
        """检测蓝色直线 - 只检测中央区域以避免H字形两边竖线干扰"""
        if image is None:
            return None, None
        
        # 转换为LAB色彩空间
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        
        # 创建蓝色掩码
        mask = cv2.inRange(lab, self.blue_lower, self.blue_upper)
        
        # 形态学操作去除噪声
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 限制检测区域为图像中央80%
        height, width = mask.shape
        center_width = int(width * 0.8)  # 中央80%宽度
        start_x = int(width * 0.1)       # 从10%位置开始
        end_x = start_x + center_width   # 到90%位置结束
        
        # 创建区域掩码，只保留中央区域
        region_mask = np.zeros_like(mask)
        region_mask[:, start_x:end_x] = 255
        
        # 应用区域限制
        mask_limited = cv2.bitwise_and(mask, region_mask)
        
        print(f"Detection region: x={start_x}-{end_x}, width={center_width}/{width}")
        
        # 使用霍夫直线变换检测线条（在限制区域内）
        lines = cv2.HoughLinesP(mask_limited, 1, np.pi/180, threshold=30, minLineLength=30, maxLineGap=15)
        
        if lines is not None and len(lines) > 0:
            # 找到最长的线条
            longest_line = None
            max_length = 0
            
            for line in lines:
                x1, y1, x2, y2 = line[0]
                length = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
                if length > max_length:
                    max_length = length
                    longest_line = line[0]
            
            if longest_line is not None:
                x1, y1, x2, y2 = longest_line
                
                # 计算线条中心点和角度
                center_x = (x1 + x2) / 2
                center_y = (y1 + y2) / 2
                
                # 计算角度（相对于水平线）
                if x2 != x1:
                    angle = math.atan2(y2 - y1, x2 - x1) * 180 / math.pi
                else:
                    angle = 90  # 垂直线
                
                # 标准化角度到-90到90度
                while angle > 90:
                    angle -= 180
                while angle < -90:
                    angle += 180
                
                return {
                    'line': longest_line,
                    'center': (center_x, center_y),
                    'angle': angle,
                    'length': max_length
                }, mask_limited  # 返回限制区域的掩码用于显示
        
        return None, mask_limited





    def control_robot_movement(self, line_info):
        """统一的机器人移动控制 - 边转向对齐边小步前进"""
        if not self.active:  # 如果任务未激活，不执行移动
            return
            
        if line_info is None:
            print("No blue line detected, searching...")
            self.gait_manager.stop()
            self.publish_status("searching")
            return
        
        # 获取线条信息
        angle = line_info['angle']
        center_x, center_y = line_info['center']
        
        print(f"Blue line: angle={angle:.1f}°, center=({center_x:.0f}, {center_y:.0f}), state={self.current_state}")
        
        if self.current_state == "approaching":
            # 检查是否到达目标位置（要求连续多帧满足以避免抖动误判）
            print(f"[DEBUG] 检查是否到达目标位置: center_y={center_y:.0f}, target_y={self.target_y}, 满足={center_y >= self.target_y}, 计数={self.target_y_counter}/{self.target_y_stable_frames}")
            if center_y >= self.target_y:
                self.target_y_counter += 1
            else:
                self.target_y_counter = 0

            if self.target_y_counter >= max(1, self.target_y_stable_frames):
                print("到达目标位置（稳定）！准备执行跨栏动作")
                self.current_state = "ready"
                self.publish_status("ready")
                return

            # 统一的转向+前进控制
            self.unified_approach_control(line_info)
            
        elif self.current_state == "ready":
            # 准备执行跨栏
            self.execute_hurdles()
    
    def unified_approach_control(self, line_info):
        """统一的接近控制 - 同时处理转向和前进"""
        angle = line_info['angle']
        center_x, center_y = line_info['center']
        
        # 计算与图像中心的水平偏差
        deviation_x = center_x - self.center_x
        
        # 计算转向输出
        yaw_output = 0
        
        print(f"Deviation_x = {deviation_x:.1f}, Angle = {angle:.1f}°") 
        
        # 基于位置偏差的左右调整 - 哪边蓝色像素多向哪边转
        if abs(deviation_x) > 10:  # 偏差阈值
            if deviation_x > 0:  # 栏杆偏右，机器人向右转跟踪
                position_correction = -(abs(deviation_x) / 20) * self.position_gain
                yaw_output += position_correction
                print(f"Hurdle right, turn right: {position_correction:.2f}")
            else:  # 栏杆偏左，机器人向左转跟踪
                position_correction = (abs(deviation_x) / 20) * self.position_gain
                yaw_output += position_correction
                print(f"Hurdle left, turn left: +{position_correction:.2f}")
        
        # 基于角度的方向调整 - 跟随栏杆方向
        if abs(angle) > 3:  # 角度阈值
            if angle > 0:  # 栏杆右端向上，机器人向右转跟随
                angle_correction = -(abs(angle) / 12) * self.angle_gain
                yaw_output += angle_correction
                print(f"Right end up, turn right: {angle_correction:.2f}")
            else:  # 栏杆左端向上，机器人向左转跟随
                angle_correction = (abs(angle) / 12) * self.angle_gain
                yaw_output += angle_correction
                print(f"Left end up, turn left: +{angle_correction:.2f}")
        
        # 限制转向幅度
        yaw_output = max(-4, min(4, yaw_output))
        print(f"Final yaw_output: {yaw_output:.2f}")

        try:
            # 执行小步前进同时转向对齐
            small_step_distance = self.step_distance * self.small_step_ratio
            
            self.gait_manager.set_step(
                self.dsp_param,          # DSP参数
                small_step_distance,     # x方向步长（小步前进）
                0,                       # y方向步长（不侧移）
                int(yaw_output),         # yaw角度（转向）
                self.walking_param,      # 步态参数
                self.arm_swap,           # 摆臂幅度
                0                        # 单步移动
            )
            
            print(f"Unified control: forward={small_step_distance:.3f}, yaw={yaw_output:.2f}")
            self.publish_status("approaching")
            
        except Exception as e:
            print(f"Error during unified control: {e}")
            rospy.logerr(f"Unified control error: {e}")
            self.publish_status("error")



    def execute_hurdles(self):
        """执行跨栏动作"""
        if self.hurdles_executed:
            return
            
        try:
            print("Executing hurdles action...")
            self.publish_status("executing_hurdles")
            
            # 停止步态管理器
            self.gait_manager.stop()
            time.sleep(0.5)
            
            # 禁用步态管理器
            self.gait_manager.disable()
            
            # 执行跨栏动作
            self.motion_manager.run_action('hurdles')
            
            # 恢复准备状态
            self.motion_manager.run_action('walk_ready')
            
            self.hurdles_executed = True
            self.task_completed = True
            self.current_state = "done"
            self.active = False
            print("Hurdles action completed!")
            self.publish_status("completed")
            common.loginfo('Hurdles action executed successfully')
            
        except Exception as e:
            print(f"Error during hurdles action: {e}")
            rospy.logerr(f"Hurdles action error: {e}")
            self.publish_status("error")

    def set_servo(self, pan, tilt):
        """设置舵机位置"""
        try:
            # 使用motion_manager控制舵机，时间200ms，舵机23为左右，舵机24为上下
            self.motion_manager.set_servos_position(200, [[23, pan], [24, tilt]])
        except Exception as e:
            rospy.logerr(f"Servo control error: {e}")

    def shutdown(self, signum, frame):
        """节点关闭回调函数"""
        with self.lock:
            self.motion_manager.run_action('stand')
            self.running = False 
            self.active = False
            self.publish_status("shutdown")
            common.loginfo('%s shutdown' % self.name)

    def run(self):
        """主循环函数"""
        print("Waiting for camera image...")
        
        # 等待图像数据
        while self.current_image is None and self.running:
            rospy.sleep(0.1)
        
        print("Camera ready, hurdles task standby...")
        self.publish_status("ready")
        
        # 窗口显示状态
        windows_created = False
        mask_window_created = False
        
        while self.running and not rospy.is_shutdown():
            try:
                # 支持运行时更新显示配置
                self.update_display_config()
                if not self.display_main_window and windows_created:
                    self._safe_destroy_window('Hurdles Tracker - Unified Control - ACTIVE')
                    windows_created = False
                if (not self.display_mask_window or not self.display_main_window) and mask_window_created:
                    self._safe_destroy_window('Blue Line Mask - ACTIVE')
                    mask_window_created = False

                if self.current_image is not None and self.active:
                    # 创建显示图像
                    display_image = self.current_image.copy()
                    
                    # 检测蓝色线条
                    line_info, mask = self.detect_blue_line(self.current_image)
                    
                    if line_info is not None:
                        # 绘制检测到的线条
                        x1, y1, x2, y2 = line_info['line']
                        cv2.line(display_image, (x1, y1), (x2, y2), (255, 0, 0), 3)
                        
                        # 绘制线条中心点
                        center_x, center_y = line_info['center']
                        cv2.circle(display_image, (int(center_x), int(center_y)), 8, (0, 255, 0), -1)
                        
                        # 显示信息
                        cv2.putText(display_image, f"Angle: {line_info['angle']:.1f}°", (10, 30), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
                        cv2.putText(display_image, f"State: {self.current_state.upper()}", (10, 60), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        cv2.putText(display_image, "HURDLES - UNIFIED CONTROL", (10, 90), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        
                        # 显示角度状态和位置状态
                        angle_status = "ALIGNED" if abs(line_info['angle']) < 5 else f"ANGLE: {line_info['angle']:.1f}°"
                        cv2.putText(display_image, angle_status, (10, 120), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0) if abs(line_info['angle']) < 5 else (0, 255, 255), 2)
                        
                        # 显示偏差信息
                        deviation_x = center_x - self.center_x
                        cv2.putText(display_image, f"Deviation: {deviation_x:.1f}px", (10, 150), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
                        
                        # 显示位置状态
                        center_x, center_y = line_info['center']
                        position_status = f"Y: {center_y:.0f}/{self.target_y} ({'REACHED' if center_y >= self.target_y else 'APPROACHING'})"
                        cv2.putText(display_image, position_status, (10, 180), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0) if center_y >= self.target_y else (255, 255, 0), 2)
                        
                        # 控制机器人移动
                        if not self.task_completed:
                            self.control_robot_movement(line_info)
                    else:
                        cv2.putText(display_image, "NO BLUE LINE DETECTED", (10, 30), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                        cv2.putText(display_image, f"State: {self.current_state.upper()}", (10, 60), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        cv2.putText(display_image, "HURDLES - UNIFIED CONTROL", (10, 90), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        if not self.task_completed:
                            self.gait_manager.stop()
                            self.publish_status("searching")
                        
                        # 显示蓝色掩码
                        if mask is not None and self.display_mask_window and self.display_main_window:
                            cv2.imshow('Blue Line Mask - ACTIVE', mask)
                            mask_window_created = True
                    
                    # 显示任务状态
                    if self.task_completed:
                        cv2.putText(display_image, "TASK COMPLETED", (10, 250), 
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)
                    
                    # 绘制目标水平线和中心线
                    cv2.line(display_image, (0, self.target_y), (self.image_width, self.target_y), (0, 255, 255), 2)
                    cv2.line(display_image, (self.center_x, 0), (self.center_x, self.image_height), (128, 128, 128), 1)
                    cv2.putText(display_image, "Target Line", (self.image_width - 150, self.target_y - 10), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
                    
                    # 绘制检测区域边界（中央80%区域）
                    start_x = int(self.image_width * 0.1)
                    end_x = int(self.image_width * 0.9)
                    cv2.line(display_image, (start_x, 0), (start_x, self.image_height), (0, 255, 0), 2)
                    cv2.line(display_image, (end_x, 0), (end_x, self.image_height), (0, 255, 0), 2)
                    cv2.putText(display_image, "Detection Zone", (start_x + 10, 50), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                    
                    # 显示图像（只在激活时）
                    if self.display_main_window:
                        cv2.imshow('Hurdles Tracker - Unified Control - ACTIVE', display_image)
                        windows_created = True

                        # 显示蓝色掩码
                        if mask is not None and self.display_mask_window:
                            cv2.imshow('Blue Line Mask - ACTIVE', mask)
                            mask_window_created = True
                        elif mask_window_created and (mask is None or not self.display_mask_window):
                            self._safe_destroy_window('Blue Line Mask - ACTIVE')
                            mask_window_created = False

                        # 检查退出键
                        key = cv2.waitKey(30) & 0xFF
                        if key == 27:  # ESC键
                            break
                    else:
                        rospy.sleep(0.05)
                        
                elif not self.active and windows_created:
                    # 任务未激活时关闭窗口
                    self._safe_destroy_window('Hurdles Tracker - Unified Control - ACTIVE')
                    if mask_window_created:
                        self._safe_destroy_window('Blue Line Mask - ACTIVE')
                        mask_window_created = False
                    windows_created = False
                    rospy.sleep(0.1)
                else:
                    rospy.sleep(0.1)
                    
            except KeyboardInterrupt:
                break
            except Exception as e:
                print(f"Error in main loop: {e}")
                rospy.sleep(0.1)
        
        # 退出前动作
        self.gait_manager.stop()
        self.set_servo(self.head_pan_init, self.head_tilt_look_down)
        self.publish_status("stopped")
        if windows_created:
            self._safe_destroy_window('Hurdles Tracker - Unified Control - ACTIVE')
        if mask_window_created:
            self._safe_destroy_window('Blue Line Mask - ACTIVE')
        rospy.signal_shutdown('shutdown')

if __name__ == '__main__':
    try:
        hurdles_tracker = HurdlesTracker('hurdles_tracker')
        hurdles_tracker.run()
    except Exception as e:
        print(f"Failed to start hurdles tracker: {e}")
        rospy.logerr(f"Hurdles tracker initialization error: {e}")
